package com.luom.fsp.manager.websocket;

import cn.hutool.json.JSON;
import cn.hutool.json.JSONUtil;
import cn.hutool.jwt.JWT;
import com.luom.fsp.model.NotificationMessage;
import com.luom.fsp.service.UserService;
import com.luom.fsp.utils.JwtUtil;
import io.swagger.models.auth.In;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import javax.annotation.Resource;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.Set;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @ClassName: WebSocketServer
 * @Description: WebSocket服务端，支持系统消息实时推送
 */
@ServerEndpoint("/ws/message/{token}")
@Component
@Slf4j
public class WebSocketServer {

    @Resource
    private UserService userService;

    // 使用双重结构的连接池：用户ID -> 会话集合
    private static final ConcurrentHashMap<Long, CopyOnWriteArraySet<WebSocketServer>> USER_CONNECTIONS = new ConcurrentHashMap<>();
    // 心跳超时时间（秒）
    private static final long HEARTBEAT_TIMEOUT = 120;

    private Session session;
    private Long userId;
    private Timer heartbeatTimer;

    @OnOpen
    public void onOpen(Session session, @PathParam("token") String token) {
        try {
            // 1. 增强Token验证
            Map<String, Object> claims = JwtUtil.parseToken(token);
            if (claims == null || !"admin".equals(claims.get("userRole"))) {
                session.close(new CloseReason(CloseReason.CloseCodes.VIOLATED_POLICY, "权限不足"));
                return;
            }

            // 2. 初始化会话
            this.session = session;
            this.userId = ((Number)claims.get("id")).longValue();

            // 3. 连接管理（同一用户最多保持3个连接）
            manageConnections();

            // 4. 启动心跳检测
            startHeartbeatCheck();

            log.info("[实时消息推送] 连接成功 | 用户:{} | 会话:{} | 当前连接数:{}",
                    userId, session.getId(), getTotalConnections());
        } catch (Exception e) {
            log.error("[实时消息推送] 连接异常: {}", e.getMessage());
        }
    }

    private void manageConnections() {
        USER_CONNECTIONS.compute(userId, (k, v) -> {
            if (v == null) v = new CopyOnWriteArraySet<>();

            // 清理无效连接
            v.removeIf(ws -> !ws.session.isOpen());

            // 控制最大连接数
            if (v.size() >= 3) {
                v.iterator().next().closeWithWarning("达到最大连接数");
            }

            v.add(this);
            return v;
        });
    }

    private void startHeartbeatCheck() {
        heartbeatTimer = new Timer(true);
        heartbeatTimer.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                try {
                    if (!session.isOpen()) {
                        cancel();
                        return;
                    }

                    long lastActive = System.currentTimeMillis() -
                            (Long) session.getUserProperties().getOrDefault("lastActive", 0L);

                    if (lastActive > HEARTBEAT_TIMEOUT * 1000) {
                        closeWithWarning("心跳超时");
                    } else {
                        session.getAsyncRemote().sendPing(ByteBuffer.wrap("ping".getBytes()));
                    }
                } catch (Exception e) {
                    cancel();
                }
            }
        }, HEARTBEAT_TIMEOUT * 1000, HEARTBEAT_TIMEOUT * 1000);
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        // 更新最后活跃时间
        session.getUserProperties().put("lastActive", System.currentTimeMillis());

        // 处理心跳响应
        if ("pong".equals(message)) {
            return;
        }

        log.info("[实时消息推送] 收到消息 | 用户:{} | 内容:{}", userId, message);
    }

    @OnClose
    public void onClose() {
        try {
            USER_CONNECTIONS.computeIfPresent(userId, (k, v) -> {
                v.remove(this);
                return v.isEmpty() ? null : v;
            });

            if (heartbeatTimer != null) {
                heartbeatTimer.cancel();
            }
            log.info("[实时消息推送] 连接关闭 | 用户:{} | 剩余连接:{}", userId, getTotalConnections());
        } catch (Exception e) {
            log.error("[实时消息推送] 关闭异常: {}", e.getMessage());
        }
    }

    private void closeWithWarning(String reason) {
        try {
            session.close(new CloseReason(
                    CloseReason.CloseCodes.VIOLATED_POLICY,
                    "连接关闭: " + reason
            ));
        } catch (IOException e) {
            log.error("[实时消息推送] 强制关闭失败: {}", e.getMessage());
        }
    }

    public static void broadcast(NotificationMessage message) {
        String json = JSONUtil.toJsonStr(message);
        USER_CONNECTIONS.forEach((userId, sessions) -> {
            sessions.forEach(ws -> {
                if (ws.session.isOpen()) {
                    synchronized (ws.session) {
                        try {
                            ws.session.getBasicRemote().sendText(json);
                        } catch (IOException e) {
                            log.warn("[实时消息推送] 消息发送失败 | 用户:{} | 错误:{}", userId, e.getMessage());
                            sessions.remove(ws);
                        }
                    }
                } else {
                    sessions.remove(ws);
                }
            });
        });
    }

    private static int getTotalConnections() {
        return USER_CONNECTIONS.values().stream()
                .mapToInt(Set::size)
                .sum();
    }
}